from core.function import Function
from core.cuda import cuda_module
from utils.common import to_variable
import utils

# 广播
class BroadcastTo(Function):

    def __init__(self, shape):
        """
        初始化 BroadcastTo 函数。

        Parameters:
            shape (tuple): 广播后的形状。
        """
        self.shape = shape

    def forward(self, x):
        """
        执行广播操作的前向传播。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 广播后的结果。
        """
        self.x_shape = x.shape
        y = cuda_module.broadcast_to(x, self.shape)
        return y

    # broadcast_to函数的反向传播会求梯度的和
    def backward(self, gy):
        """
        执行广播操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """

        # sum_to 会求x的元素之和，并将结果的形状变为shape的形状
        gx = utils.functions_collect.sum_to(gy, self.x_shape)
        return gx

    def broadcast_to(self, x):
        """
        执行广播操作。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 广播后的结果。
        """
        if x.shape == self.shape:
            return to_variable(x)
        return self(x)
